import openai
import json
from typing import Any, Callable, Optional, Union, Sequence, Dict, List
import os
from dotenv import load_dotenv
import configparser
from Server.agent.common_base_agent import CommonBaseAgent
import time
import eventlet
load_dotenv()

# 全局配置
config = configparser.ConfigParser()
config.read('config.ini')

# 模型配置字典
MODEL_CONFIGS = {
    "openai_cfg": {
        "config_name": "openai_cfg",
        "model_type": "openai_chat",
        "model_name": config['Global']['model'],
        "api_key": config['Global']['api_key'],
        "client_args": {
            "base_url": config['Global']['base_chat_url']
        }
    }
}

class TaskListSplitAgent(CommonBaseAgent):
    def __init__(
        self,
        name: str,
        model_config_name: str = 'openai_cfg',
        tools: Optional[Union[str, Dict[str, Callable], List[Callable]]] = None,
        sys_prompt: str = "You're a helpful assistant.",
        max_iters: int = 10,
        verbose: bool = False,
        isSpeak: bool = False,
        isFunctionCall: bool = False,
        openai_api_key: Optional[str] = None,
        customKey: Optional[List[dict]] = None, # [{"class_type": "AgentMsg", "handle_type": "source", "id": "custom-key-total_score-msg", "name": "total_score", "note": "\u603b\u5206", "top": 65, "type": "str", "value_type": "msg"}]
        **kwargs: Any,
    ):
        super().__init__(name, model_config_name, tools, sys_prompt, max_iters, verbose, isSpeak, isFunctionCall, openai_api_key, customKey, **kwargs)
        self.context = None
        self.memory: List[Dict] = []
        self.task_list = []
        
        # 从配置中加载模型配置
        if model_config_name not in MODEL_CONFIGS:
            raise ValueError(f"Model configuration '{model_config_name}' not found")
        
        model_config = MODEL_CONFIGS[model_config_name]
        self.model_name = model_config["model_name"]
        self.openai_api_key = openai_api_key or model_config["api_key"] or os.getenv("OPENAI_API_KEY")
        
        if not self.openai_api_key:
            raise ValueError("OpenAI API key must be provided via argument, config, or environment variable.")
        
        # 初始化 OpenAI 客户端
        client_args = model_config.get("client_args", {})
        self.client = openai.OpenAI(
            api_key=self.openai_api_key,
            **client_args
        )
        self.customKey = customKey

    def _parse_json_from_markdown(self, text: str) -> dict:
        """Parse a markdown code block containing JSON."""
        import re
        match = re.search(r"```json(.*)```", text, re.DOTALL)
        if match:
            json_str = match.group(1).strip()
            return json.loads(json_str)
        # fallback: try to parse whole text
        outputdict = json.loads(text)
        self.speak(f"requiredKeys {self.required_keys} outputdict {outputdict} customkey {self.customKey}")
        for key in self.required_keys:
            if key not in outputdict:
                return None
        return outputdict

    def speak(self, msg: Any):
        if self.verbose:
            print(f"[Agent {self.name}] {msg}")

    def setmemory(self, memory:Union[str, List[Dict]]):
        if isinstance(memory, str):
            self.memory = [{"role":"system" ,"content":memory}]
        else:
            self.memory = memory

    def reply(self, x: Optional[Union[str, Dict, List[Dict]]] = None, exe_his=None) -> dict:
        if x:
            if isinstance(x, str):
                self.memory.append({"role": "user", "content": x})
            elif isinstance(x, dict):
                if "role" in x and "content" in x:
                    self.memory.append(x)
            elif isinstance(x, list):
                for msg in x:
                    if isinstance(msg, dict) and "role" in msg and "content" in msg:
                        self.memory.append(msg)
                        
        memory = self.memory.copy()
        memory.insert(0, {"role": "system", "content": self.sys_prompt.format(name=self.name)})
        if self.context:
            memory.append({"role": "user", "content": "Here is the context you can refer to: "+self.context})
        for i in range(self.max_iters):
            self.speak(f"ITER {i+1}, STEP 1: REASONING")
            # Prepare prompt
            messages = []
            for msg in memory:
                if isinstance(msg, dict) and "role" in msg and "content" in msg and msg["content"] != "":
                    messages.append({"role": msg["role"], "content": msg["content"]})
                    # print(f"messages: {str(msg)[:500]}...{str(msg)[-500:]}")
            # Add format instruction
            format_instruction = self._format_instruction()
            messages.append({"role": "system", "content": format_instruction})
            
            self.speak("Prompt:")
            self.speak(str(messages)[-500:])
            try:
                response = self.client.chat.completions.create(
                    model=self.model_name,
                    messages=messages,
                    temperature=0.7,
                    stream=True,
                )
                full_content = []
                for chunk in response:
                    chunk_content = chunk.choices[0].delta.content
                    if chunk_content:
                        full_content.append(chunk_content)
                    if exe_his:   exe_his.update_content(''.join(full_content))
                eventlet.sleep(0)  # 让 eventlet 有机会切换
                content = ''.join(full_content)
                self.speak("Result Parsed:")
                self.speak(content)
                # Try to parse JSON from markdown
                try:
                    parsed = self._parse_json_from_markdown(content)
                    if parsed is None:
                        continue # 重新iter
                except Exception as e:
                    parsed = {"reason": content, "error": str(e)}
                self.memory.append({"role": "assistant", "content": content})
                if self.isSpeak and "speak" in parsed:
                    self.speak(parsed["speak"])
                return parsed
            except Exception as e:
                self.speak(f"Error: {e}")
                continue
        return {"error": "Max iterations reached"}

    def _format_instruction(self) -> str:
        self.content_hint = {"reason": "what you reason",'querys':'a list of querys to split the task list, like ["some words about query1", "some words about query2", "some words about query3"]'}
        self.required_keys = ["reason","querys"]
         

        return (
            f"Respond a JSON dictionary in a markdown's fenced code block as follows, you MUST make sure all your output inside the JSON format with required keys: {', '.join(self.required_keys)}, like this:\n"
            "```json\n" + json.dumps(self.content_hint, ensure_ascii=False) + "\n```\n"
            "All content outside '```json' will be ignored."
        )
